% Crank-Nicolson 法求解二维Schrodinger方程
% i \hbar \pdv{u}{t} = -\frac{\hbar^2}{2m} \pdv[2]{u}{x} + V u
% \frac{u^{(k+1)}_i - u^{(k)}_i}{\Delta t} = i\frac{\hbar}{2m} \frac{1}{2} \left( \frac{u^{(k+1)}_{i+1} - 2u^{(k+1)}_i + u^{(k+1)}_{i-1}}{(\Delta x)^2} + \frac{u^{(k)}_{i+1} - 2u^{(k)}_i + u^{(k)}_{i-1}}{(\Delta x)^2} \right)- \frac{i}{\hbar} V \frac{u^{(k+1)} + u^{(k)}}{2}
% u^{(k+1)}_i - \frac{\alpha}{2} \left( u^{(k+1)}_{i+1} - 2u^{(k+1)}_i + u^{(k+1)}_{i-1} \right) + \frac{\beta}{2} V u^{(k+1)}_i =  u^{(k)}_i + \frac{\alpha}{2} \left( u^{(k)}_{i+1} - 2u^{(k)}_i + u^{(k)}_{i-1} \right) - \frac{\beta}{2} V u^{(k)}_i
% 其中 \alpha = i\frac{\hbar}{2m} \cdot \frac{\Delta t}{(\Delta x)^2}, \quad \beta = \frac{i}{\hbar} \Delta t
% A u^{(k+1)} = B u^{(k)}, \quad
% A = I - \frac{\alpha}{2} \nabla^2 + \frac{\beta}{2} V, \quad
% B = I + \frac{\alpha}{2} \nabla^2 - \frac{\beta}{2} V
% 参考：
% https://wuli.wiki/online/CraNic.html
% https://zhuanlan.zhihu.com/p/393374195
% https://www.bilibili.com/video/BV1sq4y1V7JT
% 可能有bug
% Gitee Repo


clc
clear

L = 2;
dx = 0.05;
dt = 0.01;

x = (-L:dx:L)';
n = size(x,1);


hbar = 1;
m = 20;

u0 = zeros(n,1);
u1 = zeros(n,1);


u0 = exp(-20*(x+0.5*L).^2).*exp(i*40*x);

V = zeros(n,1);
V(x>0) = 10;

alpha = i*hbar/(2*m)*dt/(dx)^2;
beta = i/hbar*dt;

lap = spalloc(n,n,3*n);
for i = 1:n
    if i == 1 || i == n
        lap(i,i) = 1;
        continue
    end
    lap(i,i) = -2;
    lap(i,i+1) = 1;
    lap(i,i-1) = 1;
end

clear i

A = eye(n) - 1/2*alpha*lap + 1/2*beta*diag(V);;
B = eye(n) + 1/2*alpha*lap - 1/2*beta*diag(V);

for tick = 1:1000
    u1 = A\(B*u0);
    u0 = u1;

    if mod(tick,10) == 0
        clf
        hold on
        axis equal
        axis([-L,L,0,1])
        umag = real(u1).^2+imag(u1).^2;
        plot(x,umag)
        drawnow
        pause(0.1)
        disp(sum(umag))
    end
end

